import numpy as np
from DataControl.getDataFromTushare.tudata import pullDataFromTushare
from Index import utils
from scipy.stats import pearsonr
import pandas as pd
import seaborn as sns
import matplotlib.pyplot as plt


class Relation(object):
    """
    Relation 类用于计算股票与大盘指数之间的相关性，并通过热力图展示这些相关性。

    属性：
        all_tscode (ndarray): 所有股票代码的数组。
        raw_data (dict): 存储每只股票的每日涨跌幅数据的字典。
        dapan (ndarray): 从 tushare 获取的大盘指数数据。
        relation (dict): 存储各大盘指数与股票之间相关性的字典。

    方法：
        __init__(): 初始化 Relation 对象，获取所有股票代码和大盘指数数据，并计算相关性。
        getrelation(dapandata, odidata): 计算两个时间序列（大盘指数和股票）之间的相关性。
        show(): 绘制股票与大盘指数之间的相关性热力图。
    """

    def __init__(self):
        """
        初始化 Relation 对象，获取所有股票代码和大盘指数数据，并计算每只股票与各大盘指数之间的相关性。

        数据来源：
            pullDataFromTushare(): 用于从 tushare 拉取数据的工具。
        """
        # 从 tushare 拉取数据
        pull = pullDataFromTushare()
        # 获取所有股票代码
        self.all_tscode = pull.getAllTs_code()
        self.all_tscode = self.all_tscode.values
        self.raw_data = dict()

        # 获取每只股票的每日涨跌幅数据
        for i in range(self.all_tscode.shape[0]):
            ts_code = self.all_tscode[i]
            print(ts_code)
            # 获取股票的交易记录数量
            count = pull.get_count_daily(ts_code=ts_code)
            count = count.iloc[0, 0]
            # 如果记录数量不大于100，则跳过
            if count <= 100: continue
            # 获取股票的每日涨跌幅数据
            pct_chg = pull.getallindex_daily_from_database(ts_code=ts_code)[['trade_date', 'pct_chg']].values
            self.raw_data[ts_code] = pct_chg

        # 获取大盘指数数据
        self.dapan = pull.get_dapanfromtushare().values
        self.relation = dict()

        # 计算每个大盘指数与每只股票之间的相关性
        for i in range(len(self.dapan)):
            self.relation[self.dapan[i][0]] = dict()
            if self.dapan[i][0] not in self.raw_data: continue
            for j in range(len(self.all_tscode)):
                dapanst, odist = self.dapan[i][0], self.all_tscode[j]
                if dapanst in self.raw_data and odist in self.raw_data:
                    self.relation[dapanst][odist], _ = self.getrelation(self.raw_data[dapanst], self.raw_data[odist])

    def getrelation(self, dapandata, odidata):
        """
        计算两个时间序列（大盘指数和股票）之间的相关性。

        参数：
            dapandata (ndarray): 大盘指数的每日涨跌幅数据。
            odidata (ndarray): 股票的每日涨跌幅数据。

        返回：
            tuple: (相关系数, p-value)，如果序列长度不足以计算相关性，返回 (None, None)。
        """
        dapanidx = dapandata.shape[0] - 1
        cnt = 0
        dapanlst = []
        odilst = []
        # 从后向前匹配数据
        for i in range(odidata.shape[0] - 1, odidata.shape[0] - 101, -1):
            while dapanidx >= 0 and dapandata[dapanidx][0] != odidata[i][0]:
                print(dapandata[dapanidx][0], odidata[i][0])
                dapanidx -= 1
            if dapanidx == -1: break
            cnt += 1
            dapanlst.append(dapandata[dapanidx][1])
            odilst.append(odidata[i][1])
            dapanidx -= 1

        # 如果匹配到的点数等于100，计算皮尔逊相关系数
        if cnt == 100:
            return pearsonr(dapanlst, odilst)
        else:
            return None, None

    def show(self):
        """
        绘制股票与大盘指数之间的相关性热力图。
        """
        # 将相关性数据转换为 DataFrame
        df = pd.DataFrame(self.relation)
        # 设置绘图大小
        plt.figure(figsize=(14, 10))
        # 创建热力图
        sns.heatmap(df, cmap='coolwarm', linewidths=.5)
        # 添加标题和标签
        plt.title('Correlation Heatmap between Stocks and Indices')
        plt.xlabel('Stocks')
        plt.ylabel('Indices')
        # 显示图形
        plt.show()


if __name__ == "__main__":
    relation = Relation()
    relation.show()
    #print(relation.relation['000001.SH']['000001.SH'])
    #print(relation.relation)